import torch
import torch.nn as nn


def activate_function(x): return nn.SiLU(inplace=True)(x)


class NeuralBlock(nn.Module):  # 神经网络块
    def __init__(self, input_size, output_size):
        super(NeuralBlock, self).__init__()
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, x):
        return activate_function(self.fc(x))


class AutoDecoderNet(nn.Module):  # 自动解码器神经网络
    def __init__(self, coordinate_size: int, hidden_size: int, output_size: int, depth: int,  latent_vectors_num: int, latent_size: int):
        super(AutoDecoderNet, self).__init__()
        self.geometry_vectors_scale = 0.01
        self.latent_vectors_num = latent_vectors_num
        geometry_size = 3
        self.layers = nn.ModuleList()  # 使用 ModuleList 来存储所有的层
        self.layers.append(nn.Linear(coordinate_size +
                           latent_size + geometry_size, hidden_size))  # 添加输入层

        for i in range(depth - 1):  # 添加神经网络块
            self.layers.append(NeuralBlock(
                hidden_size + latent_size + geometry_size, hidden_size))

        self.layers.append(nn.Linear(hidden_size, output_size))  # 添加输出层

        self.latent_vectors = nn.Parameter(
            torch.FloatTensor(latent_vectors_num, latent_size))
        nn.init.xavier_normal_(self.latent_vectors, 0.01)

        self.geometry_vectors = nn.Parameter(
            torch.zeros(latent_vectors_num, geometry_size), requires_grad=False)

    def forward(self, x, lv, gv):
        for layer in self.layers[:-1]:
            x = self.concatenate_vectors_to_x(x, lv)
            x = self.concatenate_vectors_to_x(x, gv)  # 将隐向量、几何向量与x拼接
            x = layer(x)
        x = self.layers[-1](x)
        return x

    def forward_for_interpolation(self, x, latent_vector_index1, latent_vector_index2, alpha):
        lv = self.latent_vector_interpolation(
            latent_vector_index1, latent_vector_index2, alpha)
        gv = self.geometry_vector_interpolation(
            latent_vector_index1, latent_vector_index2, alpha)
        return self.forward(x, lv, gv)

    def forward_with_latent_vector_index(self, x, latent_vector_index):
        lv = self.get_latent_vectors(latent_vector_index)
        gv = self.get_geometry_vectors(latent_vector_index)
        return self.forward(x, lv, gv)

    def concatenate_vectors_to_x(self, x, v):
        """
        将潜在向量 v 拼接到输入特征 x 上。
        
        参数:
            x: 输入特征张量，形状为 (batch_size, *spatial_dims, feature_dim)。
            v: 潜在向量，形状为 (batch_size, latent_dim)。
        
        返回:
            拼接后的张量。
        """
        # 获取 x 的维度信息
        batch_size = x.shape[0]
        spatial_dims = x.shape[1:-1]  # 空间维度（如高度、宽度等）

        # 扩展潜在向量 v 的维度以匹配 x 的空间维度
        v_expanded = v.view(batch_size, *((1,) * len(spatial_dims)), -1)  # 扩展空间维度
        v_expanded = v_expanded.expand(batch_size, *spatial_dims, -1)  # 复制到匹配的空间维度

        # 在最后一个维度上拼接 x 和 v
        return torch.cat((x, v_expanded), dim=-1)
        # 添加潜在向量到输入特征
        #if x.shape.__len__() == 5:
        #    v = v.unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(
        #        1, x.shape[1], x.shape[2], x.shape[3], 1)
        #    return torch.cat((x, v), dim=4)
        #elif x.shape.__len__() == 4:
        #    v = v.unsqueeze(1).unsqueeze(1).repeat(
        #        1, x.shape[1], x.shape[2], 1)
        #    return torch.cat((x, v), dim=3)
        #else:
        #    v = v.unsqueeze(1).repeat(1, x.shape[1], 1)
        #    return torch.cat((x, v), dim=2)

    def latent_vector_interpolation(self, latent_vector_index1, latent_vector_index2, alpha):
        return self.latent_vectors[latent_vector_index1] * (1 - alpha) + self.latent_vectors[latent_vector_index2] * alpha

    def geometry_vector_interpolation(self, latent_vector_index1, latent_vector_index2, alpha):
        return self.geometry_vectors[latent_vector_index1] * (1 - alpha) + self.geometry_vectors[latent_vector_index2] * alpha

    def get_latent_vectors(self, latent_vector_index):
        return self.latent_vectors[latent_vector_index]

    def get_geometry_vectors(self, latent_vector_index):
        return self.geometry_vectors[latent_vector_index]

    def set_geometry_vectors(self, dataset):
        for i in range(self.latent_vectors_num):
            self.geometry_vectors[i] = dataset.get_geometry_vector(
                i).cuda()* self.geometry_vectors_scale
            